from classification.vggModel import model
from keras.optimizers import SGD
from keras.preprocessing.image import ImageDataGenerator
import os
from collections import Counter
from config import Config

if Config.ISGPU:
    from tensorflow.compat.v1 import ConfigProto
    from tensorflow.compat.v1 import InteractiveSession

    os.environ["CUDA_VISIBLE_DEVICES"] = Config.GPU
    config = ConfigProto()
    config.gpu_options.allow_growth = True
    session = InteractiveSession(config=config)

# 图片生成器
train_datagen = ImageDataGenerator(
    rotation_range=40,  # 随机旋转度数
    width_shift_range=0.2,  # 随机水平平移
    height_shift_range=0.2,  # 随机竖直平移
    rescale=1 / 255,  # 数据归一化
    shear_range=20,  # 随机错切变换
    zoom_range=0.2,  # 随机放大
    horizontal_flip=True,  # 水平翻转
    fill_mode='nearest',  # 填充方式
)
test_datagen = ImageDataGenerator(
    rescale=1 / 255,  # 数据归一化
)

# 生成训练数据
train_generator = train_datagen.flow_from_directory(
    Config.CLASS_TRAIN_DATA,
    target_size=Config.CLASS_TARGET_SIZE,
    batch_size=Config.CLASS_BATCH_SIZE,
)
class_weight = {}
counter = Counter(train_generator.classes)
max_val = float(max(counter.values()))
class_weights = {class_id: max_val / num_images for class_id, num_images in counter.items()}
print("class_weights for samples:", class_weights)

# 测试数据
test_generator = test_datagen.flow_from_directory(
    Config.CLASS_TEST_DATA,
    target_size=Config.CLASS_TARGET_SIZE,
    batch_size=Config.CLASS_BATCH_SIZE,
)
model.compile(optimizer=SGD(lr=Config.CLASS_LR, momentum=0.9), loss='categorical_crossentropy', metrics=['accuracy'])

# 训练
model.fit_generator(train_generator, steps_per_epoch=len(train_generator), epochs=Config.CLASS_EPOCHS,
                    validation_data=test_generator,
                    validation_steps=len(test_generator), class_weight=class_weights)

# 模型保存
model.save(Config.CLASS_MODEL_HOME)
